package org.apache.lucene.ocean.database;

import java.io.IOException;

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.ocean.Index.IndexSnapshot;
import org.apache.lucene.ocean.util.Timeout;
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.FieldSortedHitQueue;
import org.apache.lucene.search.Filter;
import org.apache.lucene.search.HitCollector;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.Searchable;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.Weight;
import org.apache.lucene.util.OpenBitSet;
import org.apache.lucene.util.PriorityQueue;

public class SearchKey extends CacheKey<SearchResults> {
  public Weight weight;
  public Sort sort;
  public int num;
  public int start;
  public Filter filter;
  
  public SearchResults createObject(IndexSnapshot indexSnapshot, Timeout timeout) throws Exception {
    Searchable searchable = indexSnapshot.getSearcher();
    OpenBitSet bitSet = new OpenBitSet(searchable.maxDoc());
    TopDocCollector tdc = null;
    if (sort == null) {
      tdc = new TopDocCollector(start + num, bitSet);
      searchable.search(weight, filter, tdc);
    } else {
      IndexReader reader = indexSnapshot.getIndexReader();
      tdc = new TopFieldDocCollector(reader, sort, num+start, bitSet);
      searchable.search(weight, filter, tdc);
    }
    int totalHits = tdc.getTotalHits();
    ScoreDoc[] scoreDocs = tdc.topDocs().scoreDocs;
    float maxScore = tdc.topDocs().getMaxScore();
    return new SearchResults(scoreDocs, bitSet, totalHits, maxScore);
  }
  
  public class TopFieldDocCollector extends TopDocCollector {
    private FieldDoc reusableFD;

    public TopFieldDocCollector(IndexReader reader, Sort sort, int numHits, OpenBitSet bitSet) throws IOException {
      super(numHits, new FieldSortedHitQueue(reader, sort.getSort(), numHits), bitSet);
    }

    public void collect(int doc, float score) {
      if (score > 0.0f) {
        totalHits++;
        if (reusableFD == null)
          reusableFD = new FieldDoc(doc, score);
        else {
          // Whereas TopDocCollector can skip this if the
          // score is not competitive, we cannot because the
          // comparators in the FieldSortedHitQueue.lessThan
          // aren't in general congruent with "higher score
          // wins"
          reusableFD.score = score;
          reusableFD.doc = doc;
        }
        reusableFD = (FieldDoc) hq.insertWithOverflow(reusableFD);
      }
    }

    // javadoc inherited
    public TopDocs topDocs() {
      FieldSortedHitQueue fshq = (FieldSortedHitQueue) hq;
      ScoreDoc[] scoreDocs = new ScoreDoc[fshq.size()];
      for (int i = fshq.size() - 1; i >= 0; i--)
        // put docs in array
        scoreDocs[i] = fshq.fillFields((FieldDoc) fshq.pop());

      return new TopFieldDocs(totalHits, scoreDocs, fshq.getFields(), fshq.getMaxScore());
    }
  }

  public static class TopFieldDocs extends TopDocs {
    public SortField[] fields;

    public TopFieldDocs(int totalHits, ScoreDoc[] scoreDocs, SortField[] fields, float maxScore) {
      super(totalHits, scoreDocs, maxScore);
      this.fields = fields;
    }
  }

  public class TopDocCollector extends HitCollector {
    private ScoreDoc reusableSD;
    int totalHits;
    PriorityQueue hq;
    OpenBitSet bitSet;

    public TopDocCollector(int numHits, OpenBitSet bitSet) {
      this(numHits, new HitQueue(numHits), bitSet);
    }

    TopDocCollector(int numHits, PriorityQueue hq, OpenBitSet bitSet) {
      this.hq = hq;
      this.bitSet = bitSet;
    }

    public void collect(int doc, float score) {
      if (score > 0.0f) {
        bitSet.fastSet(doc);
        totalHits++;
        if (reusableSD == null) {
          reusableSD = new ScoreDoc(doc, score);
        } else if (score >= reusableSD.score) {
          // reusableSD holds the last "rejected" entry, so, if
          // this new score is not better than that, there's no
          // need to try inserting it
          reusableSD.doc = doc;
          reusableSD.score = score;
        } else {
          return;
        }
        reusableSD = (ScoreDoc) hq.insertWithOverflow(reusableSD);
      }
    }

    /** The total number of documents that matched this query. */
    public int getTotalHits() {
      return totalHits;
    }

    /** The top-scoring hits. */
    public TopDocs topDocs() {
      ScoreDoc[] scoreDocs = new ScoreDoc[hq.size()];
      for (int i = hq.size() - 1; i >= 0; i--)
        // put docs in array
        scoreDocs[i] = (ScoreDoc) hq.pop();

      float maxScore = (totalHits == 0) ? Float.NEGATIVE_INFINITY : scoreDocs[0].score;

      return new TopDocs(totalHits, scoreDocs, maxScore);
    }
  }

  public static class HitQueue extends PriorityQueue {
    public HitQueue(int size) {
      initialize(size);
    }

    protected final boolean lessThan(Object a, Object b) {
      ScoreDoc hitA = (ScoreDoc) a;
      ScoreDoc hitB = (ScoreDoc) b;
      if (hitA.score == hitB.score)
        return hitA.doc > hitB.doc;
      else
        return hitA.score < hitB.score;
    }
  }
}
